namespace MoreLinq.Experimental;

using System;
using System.Collections;
using System.Collections.Generic;
using System.Runtime.ExceptionServices;

static partial class ExperimentalEnumerable
{
    /// <summary>
    /// Creates a sequence that lazily caches the source as it is iterated
    /// for the first time, reusing the cache thereafter for future
    /// re-iterations. If the source is already cached or buffered then it
    /// is returned verbatim.
    /// </summary>
    /// <typeparam name="T">
    /// Type of elements in <paramref name="source"/>.</typeparam>
    /// <param name="source">The source sequence.</param>
    /// <returns>
    /// Returns a sequence that corresponds to a cached version of the input
    /// sequence.
    /// </returns>
    /// <exception cref="ArgumentNullException">
    /// <paramref name="source"/> is <see langword="null"/>.
    /// </exception>
    /// <remarks>
    /// The returned <see cref="IEnumerable{T}"/> will cache items from
    /// <paramref name="source"/> in a thread-safe manner. Each thread can
    /// call its <see cref="IEnumerable{T}.GetEnumerator"/> to acquire an
    /// iterator  but the same iterator should not be used simultaneously
    /// from multiple threads. The sequence supplied in <paramref
    /// name="source"/> is not expected to be thread-safe but it is required
    /// to be thread-agnostic because different threads (though never
    /// simultaneously) may iterate over the sequence.
    /// </remarks>

    public static IEnumerable<T> Memoize<T>(this IEnumerable<T> source) =>
        source switch
        {
            null => throw new ArgumentNullException(nameof(source)),
            ICollection<T> => source,
            IReadOnlyCollection<T> => source,
            MemoizedEnumerable<T> => source,
#pragma warning disable IDE0306 // Simplify collection initialization (would change semantics)
            _ => new MemoizedEnumerable<T>(source),
#pragma warning restore IDE0306 // Simplify collection initialization
        };
}

sealed class MemoizedEnumerable<T>(IEnumerable<T> sequence) : IEnumerable<T>, IDisposable
{
    List<T>? cache;
    readonly Lock locker = new();
    readonly IEnumerable<T> source = sequence ?? throw new ArgumentNullException(nameof(sequence));
    IEnumerator<T>? sourceEnumerator;
    int? errorIndex;
    ExceptionDispatchInfo? error;

    public IEnumerator<T> GetEnumerator()
    {
            if (this.cache == null)
            {
                lock (this.locker)
                {
#pragma warning disable CA1508 // Avoid dead conditional code
                    if (this.cache == null)
#pragma warning restore CA1508 // Avoid dead conditional code
                    {
                        this.error?.Throw();

                        try
                        {
                            var cache = new List<T>(); // for exception safety, allocate then...
                            this.sourceEnumerator = this.source.GetEnumerator(); // (because this can fail)
                            this.cache = cache; // ...commit to state
                        }
                        catch (Exception ex)
                        {
                            this.error = ExceptionDispatchInfo.Capture(ex);
                            throw;
                        }
                    }
                }
            }

            return _(); IEnumerator<T> _()
            {
                var index = 0;

                while (true)
                {
                    T current;
                    lock (this.locker)
                    {
                        if (this.cache == null) // Cache disposed during iteration?
                            throw new ObjectDisposedException(nameof(MemoizedEnumerable<T>));

                        if (index >= this.cache.Count)
                        {
                            if (index == this.errorIndex)
                                Assume.NotNull(this.error).Throw();

                            if (this.sourceEnumerator == null)
                                break;

                            bool moved;
                            try
                            {
                                moved = this.sourceEnumerator.MoveNext();
                            }
                            catch (Exception ex)
                            {
                                this.error = ExceptionDispatchInfo.Capture(ex);
                                this.errorIndex = index;
                                this.sourceEnumerator.Dispose();
                                this.sourceEnumerator = null;
                                throw;
                            }

                            if (moved)
                            {
                                this.cache.Add(this.sourceEnumerator.Current);
                            }
                            else
                            {
                                this.sourceEnumerator.Dispose();
                                this.sourceEnumerator = null;
                                break;
                            }
                        }

                        current = this.cache[index];
                    }

                    yield return current;
                    index++;
                }
            }
        }

    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();

    public void Dispose()
    {
            lock (this.locker)
            {
                this.error = null;
                this.cache = null;
                this.errorIndex = null;
                this.sourceEnumerator?.Dispose();
                this.sourceEnumerator = null;
            }
        }
}
